package com.liuqi.openai.core.merge;

import com.alibaba.fastjson.JSONObject;
import com.liuqi.openai.core.chat.ChatCompletionChoice;
import com.liuqi.openai.core.chat.ChatCompletionResponse;
import com.liuqi.openai.core.chat.Delta;
import com.liuqi.openai.core.chat.Role;
import com.liuqi.openai.core.chat.ToolCall;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.extern.slf4j.Slf4j;

/**
 * @author liuqi
 * @date 2025/8/6
 **/
@Slf4j
public class MergeChunkChatCompletion<Response extends ChatCompletionResponse> extends MergeChunk<Response> {

    @Override
    public void merge(ChatCompletionResponse response) {
        // 设置id、created、model、object、usage
        setCommon(response.getId(), response.getCreated(), response.getModel(), response.getObject(), response.getUsage());

        List<ChatCompletionChoice> choices = response.getChoices();
        if (choices != null && !choices.isEmpty()) {
            choices.forEach(choice -> {
                Integer index = choice.getIndex();

                Delta delta = choice.getDelta();
                if (delta != null) {
                    if (delta.getRole() != null) {
                        appendRole(index, delta.getRole().toString());
                    }
                    if (delta.getReasoningContent() != null) {
                        appendReasoningContent(index, delta.getReasoningContent());
                    }
                    if (delta.getContent() != null) {
                        appendContent(index, delta.getContent());
                    }
                    if (delta.getToolCalls() != null && !delta.getToolCalls().isEmpty()) {
                        delta.getToolCalls().forEach(toolCall -> {
                            appendToolCalls(index, toolCall);
                        });
                    }
                }

                setFinishReason(index, choice.getFinishReason());
            });
        }
    }

    @Override
    public Response getResult() {
        return (Response) ChatCompletionResponse.builder()
                .id(id)
                .created(created)
                .model(model)
                .object(object)
                .usage(usage)
                .choices(buildChatCompletionChoice())
                .build();
    }

    private List<ChatCompletionChoice> buildChatCompletionChoice() {
        List<ChatCompletionChoice> choices;
        int size;
        if (toolCallsBuff != null) {
            choices = new ArrayList<>(size = toolCallsBuff.size());
        } else if (contentBuff != null) {
            choices = new ArrayList<>(size = contentBuff.size());
        } else {
            return null;
        }
        for (int i = 0; i < size; i++) {
            // 思考内容
            choices.add(ChatCompletionChoice.builder()
                    .index(i)
                    .finishReason(getFinishReason(i))
                    .delta(buildDelta(i))
                    .build());
        }

        return choices;
    }

    private Delta buildDelta(int index) {
        return Delta.builder()
                .role(Role.from(getRole(index)))
                .reasoningContent(getReasoningContent(index))
                .content(getContent(index))
                .toolCalls(buildToolCalls(index))
                .build();
    }

    private List<ToolCall> buildToolCalls(int index) {
        if (toolCallsBuff == null) {
            return null;
        }
        List<ToolCall> toolCalls = toolCallsBuff.get(index);
        if (toolCalls == null || toolCallsBuff.isEmpty()) {
            return null;
        }

        // 由index进行归类
        Map<Integer, InternalToolCall> internalToolCalls = new HashMap<>(1);
        toolCalls.forEach(toolCall -> {
            Integer i = toolCall.getIndex();
            InternalToolCall internalToolCall = internalToolCalls.get(i);
            if (internalToolCall == null) {
                internalToolCalls.put(i, internalToolCall = new InternalToolCall());
            }
            ToolCall.Function function = toolCall.getFunction();
            internalToolCall.merge(
                    toolCall.getIndex(),
                    toolCall.getId(),
                    toolCall.getType(),
                    function == null ? null : function.getName(),
                    function == null ? null : function.getArguments()
            );
        });

        // 构建ToolCall
        List<ToolCall> resultToolCalls = new ArrayList<>(internalToolCalls.size());
        for (int i = 0; i < internalToolCalls.size(); i++) {
            InternalToolCall internalToolCall = internalToolCalls.get(i);

            // 计算 Function.arguments
            String arguments = null;
            if (internalToolCall.argumentsBuf.isEmpty()) {
                // 如果参数为空, 强制赋值为空json
                arguments = "{}";
            } else {
                // 有些模型最后一个响应块可能会包含一个完整的arguments json串
                String str = internalToolCall.argumentsBuf.get(internalToolCall.argumentsBuf.size() - 1);
                try {
                    // 是否能够正常进行反序列化, 则是一个完整的arguments json串
                    JSONObject.parseObject(str);
                    arguments = str;
                } catch (Throwable t) {
                }

                // 将arguments进行拼接
                if (arguments == null) {
                    StringBuilder buf = new StringBuilder();
                    internalToolCall.argumentsBuf.forEach(argument -> {
                        buf.append(argument);
                    });
                    arguments = buf.toString();
                    try {
                        // 校验一下能否正常序列化
                        JSONObject.parseObject(arguments);
                    } catch (Throwable t) {
                        log.error("AI responded to an incorrect tool call, input parameter: " + arguments);
                    }
                }
            }

            resultToolCalls.add(ToolCall.builder()
                    .index(i)
                    .id(internalToolCall.id)
                    .type(internalToolCall.type)
                    .function(
                            ToolCall.Function.builder()
                                    .name(internalToolCall.nameBuf.toString())
                                    .arguments(arguments)
                                    .build()
                    )
                    .build());
        }

        return resultToolCalls;
    }

    private class InternalToolCall {
        int index;
        String id;
        String type;
        StringBuilder nameBuf = new StringBuilder(6);
        List<String> argumentsBuf = new ArrayList<>(6);

        void merge(Integer index, String id, String type, String name, String arguments) {
            if (index != null) {
                this.index = index;
            }
            if (id != null && !id.isEmpty()) {
                this.id = id;
            }
            if (type != null&& !type.isEmpty()) {
                this.type = type;
            }
            if (name != null) {
                this.nameBuf.append(name);
            }
            if (arguments != null) {
                this.argumentsBuf.add(arguments);
            }
        }
    }

}
